import sys
import cv2
sys.path.append("..")
import numpy as np
from u2net import detect
import time

print('loading model......')
start1 = time.time()
model_u2net = detect.load_model(model_name="u2net")
model_u2netp = detect.load_model(model_name="u2netp")
end1 = time.time()
print(end1 - start1, '\n')



def get_target(img):
    left_point = (-1, -1)
    right_point = (-1, -1)
    for high in reversed(range(200, 310)):
        for wide in range(10, 160):
            if img[high, wide] >= 0.8:
                left_point = (high-8, wide+1)
                break
        else:
            continue
        break
    for high in reversed(range(200, 310)):
        for wide in range(160, 310):
            if img[high, wide] >= 0.8:
                right_point = (high-8, wide+1)
                break
        else:
            continue
        break
    min_value_left = img[left_point]
    for high in range(left_point[0]-3, left_point[0]+3):
        for wide in range(left_point[1]-3, left_point[1]+3):
            if img[high][wide] <= min_value_left:
                min_left = (high, wide)
    min_value_right = img[right_point]
    for high in range(right_point[0] - 3, right_point[0] + 3):
        for wide in range(right_point[1] - 3, right_point[1] + 3):
            if img[high][wide] <= min_value_right:
                min_right = (high, wide)
    return min_left, min_right


def center_line(left_point, right_point):
    x1, y1 = left_point[1], left_point[0]
    x2, y2 = right_point[1], right_point[0]
    x3 = (x1 + x2) / 2
    y3 = (y1 + y2) / 2
    d = int(((y2 - y1)**2 + (x2 - x1)**2)**0.5)
    if y1 == y2:
        return (int(x3), int(y3)), (int(x3), int(y3 - 1.74*d))
    else:
        k2 = (x1 - x2) / (y2 - y1)
        if y2 < y1:
            x4 = x3 - (((1.74*d)**2)/(1+k2**2))**0.5
        else:
            x4 = x3 + (((1.74 * d) ** 2) / (1 + k2 ** 2)) ** 0.5
        y4 = y3 + k2 * (x4 - x3)

        return (int(x3), int(y3)), (int(x4), int(y4))

def remove(model_name="u2net"):
    model = model_u2net
    if model == "u2netp":
        model = model_u2netp
    cap = cv2.VideoCapture('test.mp4')
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter('output.avi', fourcc, 20.0, (320, 320))
    ret, frame = cap.read()
    target_point = [(240, 105), (240, 215)]
    start2 = time.time()
    while ret:
        #测试图片时
        #frame = cv2.imread('test.jpg')

        frame = cv2.resize(frame, dsize=(320, 320))
        roi = detect.predict(model, frame)
        detect_point = get_target(roi)
        point1, point2 = center_line(detect_point[0], detect_point[1])
        print(point1, point2)
        if abs(detect_point[0][0] - target_point[0][0]) < 10 \
            and abs(detect_point[0][1] - target_point[0][1]) < 10 \
            and abs(detect_point[1][1] - target_point[1][1]) < 10 \
            and abs(detect_point[1][0] - target_point[1][0]) < 10:
            cv2.circle(frame, (target_point[0][1], target_point[0][0]), 5, (0, 255, 0), -1)
            cv2.circle(frame, (target_point[1][1], target_point[1][0]), 5, (0, 255, 0), -1)
            cv2.line(frame, point1, point2, (0, 255, 0), 5)
        else:
            cv2.circle(frame, (detect_point[0][1], detect_point[0][0]), 5, (0, 0, 255), -1)
            cv2.circle(frame, (detect_point[1][1], detect_point[1][0]), 5, (0, 0, 255), -1)
            cv2.circle(frame, (target_point[0][1], target_point[0][0]), 5, (255, 0, 0), -1)
            cv2.circle(frame, (target_point[1][1], target_point[1][0]), 5, (255, 0, 0), -1)
            cv2.line(frame, point1, point2, (0, 0, 255), 5)

        #显示图片
        cv2.imshow('1111', frame)
        #cv2.waitKey(0)
        #写视频
        out.write(frame)
        if cv2.waitKey(30) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
            break
        ret, frame = cap.read()
    end2 = time.time()
    print(end2 - start2)
    cap.release()
if __name__ == "__main__":
    remove()

